package com.feidee.fd.sml.algorithm.feature

import com.feidee.fd.sml.algorithm.component.feature.{StandardScaleEncoder, StandardScaleEncoderParam}
import com.feidee.fd.sml.algorithm.util.{TestingDataGenerator, ToolClass}
import org.apache.spark.ml.linalg.Vectors
import org.scalatest.FunSuite

/**
  * @Author tangjinyuan
  * @Date 2019/3/28 14:32
  * @Description
  * @Reviewer
  */
class StandardScaleEncoderSuite extends FunSuite {
  val paramStr: String =
    """
      |{
      |	"inputCol": "features",
      |	"outputCol": "scaledFeatures",
      |	"withMean": false,
      |	"withStd": true
      |}
    """.stripMargin

  val model = new StandardScaleEncoder();
  val param: StandardScaleEncoderParam = model.parseParam(new ToolClass().encrypt(paramStr))

  test("model parameter") {
    assert("features".equals(param.inputCol) && "scaledFeatures".equals(param.outputCol) && !param.withMean&& param.withStd)
  }


  val dataFrame = TestingDataGenerator.spark.createDataFrame(Seq(
    (0, Vectors.dense(1.0, 0.5, -1.0)),
    (1, Vectors.dense(2.0, 1.0, 1.0)),
    (2, Vectors.dense(4.0, 10.0, 2.0))
  )).toDF("id", "features")

  val compareSeq = Seq(
    (Vectors.dense(0.6546536707079772,0.09352195295828246,-0.6546536707079771),""),
  (Vectors.dense(1.3093073414159544,0.18704390591656492,0.6546536707079771),""),
  (Vectors.dense(2.618614682831909,1.8704390591656492,1.3093073414159542),"")

  )

  test("model transformation") {
    var compareData = TestingDataGenerator.spark.createDataFrame(compareSeq).toDF("scaledFeatures", "b")
    println("compareData:")
    compareData.select("scaledFeatures").show(false)
    println(compareData.select("scaledFeatures").schema)
    val res = model.train(param, dataFrame).transform(dataFrame)
    println("res:")
    res.select("scaledFeatures").show(false)
    println(res.select("scaledFeatures").schema)

    assert(compareData.select("scaledFeatures").except(res.select("scaledFeatures")).count() == 0)
  }


}

